//
// Created by QU on 24-9-21.
//

#include "13tlv_echo_server.h"

Session::Session(ip::tcp::socket&& socket, Server& server)
    : m_server(server), m_socket(std::move(socket)), m_uuid(boost::uuids::random_generator()()),
      m_recv_buffer(Node::MAX_HEADER_SIZE)
{
    std::cout << "Session::Session(): " << GetUuid() << std::endl;
}

Session::~Session()
{
    std::cout << "Session::~Session(): " << GetUuid() << std::endl;
}

void Session::StartSession()
{
    // std::print("Session::StartSession(): {}\n", GetUuid());
    auto self = shared_from_this(); // 确保 Session 被 shared_ptr 管理
    StartAsyncRead();
}

void Session::StartAsyncRead()
{
    m_recv_buffer.resize(MIN_BUFFER_SIZE);
    m_recv_buffer_curr_len = 0;
    m_header_received_flag = false;
    m_temp_payload_size = 0;
    m_buffer_need_shrink_counter = 0;

    ContinueRead();
}

void Session::ContinueRead()
{
    auto self = shared_from_this();
    m_socket.async_read_some(
        asio::buffer(m_recv_buffer.data() + m_recv_buffer_curr_len,
                     m_recv_buffer.size() - m_recv_buffer_curr_len),
        [self, this](const boost::system::error_code& ec, const std::size_t bytes_transferred)
        {
            HandleRead(ec, bytes_transferred);
        }
    );
}

void Session::HandleRead(const boost::system::error_code& ec, std::size_t bytes_transferred)
{
    // std::this_thread::sleep_for(std::chrono::milliseconds(2));

    if (ec)
    {
        HandleError(ec);
        return;
    }

    std::print("Session::handleRead(): Received: {} bytes.\n", bytes_transferred);
    m_recv_buffer_curr_len += bytes_transferred;

    while (m_recv_buffer_curr_len > 0)
    {
        // 开始处理头部
        if (!m_header_received_flag)
        {
            // 如果收到的数据足够一个头
            if (m_recv_buffer_curr_len >= Node::MAX_HEADER_SIZE)
            {
                const uint16_t payload_size = boost::endian::big_to_native(
                    *reinterpret_cast<uint16_t*>(&m_recv_buffer[2]));

                std::print("Session::handleRead(): Header parse: payload_size: {} bytes.\n", payload_size);

                if (payload_size > Node::MAX_PAYLOAD_SIZE)
                {
                    std::cerr << "Invalid payload size: " << payload_size << std::endl;
                    HandleError(ec, "Invalid payload size");
                    return;
                }

                // 这个包的头部, 存起来.
                m_temp_payload_size = payload_size;
                // 收到了这个包的头, 确定是否需要增加buffer的大小.
                const size_t required_size = m_temp_payload_size + Node::MAX_HEADER_SIZE;
                if (required_size > m_recv_buffer.size())
                {
                    if (!AdjustBufferSize(required_size))
                    {
                        // buffer扩容处理失败
                        HandleError(ec, "Failed to adjust buffer size");
                        return;
                    }
                }

                // 标记头部处理完成
                m_header_received_flag = true;
            }
            else
            {
                break; // 需要更多数据来完成头部
            }
        }
        // 处理payload
        else
        {
            const size_t total_message_size = m_temp_payload_size + Node::MAX_HEADER_SIZE;
            if (m_recv_buffer_curr_len >= total_message_size)
            {
                const auto recv_node = BufferProcessor(m_recv_buffer, total_message_size);
                if (!recv_node)
                {
                    std::cerr << "Failed to process message." << std::endl;
                    HandleError(ec, "Failed to process message");
                    return;
                }

                // 移动剩余数据到缓冲区开头
                const size_t remaining_data = m_recv_buffer_curr_len - total_message_size;
                if (remaining_data > 0)
                {
                    std::memmove(m_recv_buffer.data(), m_recv_buffer.data() + total_message_size, remaining_data);
                }
                m_recv_buffer_curr_len = remaining_data;

                // 重置状态
                m_header_received_flag = false;
                m_temp_payload_size = 0;
            }
            else
            {
                break; // 需要更多数据来完成消息
            }
        }
    }

    AdjustBufferSize();
    ContinueRead();

    std::cout << "\t\t\t\t\t\t\t\t\t\tBuffer usage: " << m_recv_buffer_curr_len << "/" << m_recv_buffer.size() <<
        " bytes." << std::endl;
}

bool Session::AdjustBufferSize(const size_t required_size)
{
    size_t new_size = m_recv_buffer.size();
    bool need_resize = false;

    if (required_size > 0 && required_size > new_size)
    {
        // 需要立即扩容buffer
        new_size = required_size;
        need_resize = true;
    }
    else if (m_recv_buffer_curr_len > new_size * EXPAND_THRESHOLD)
    {
        new_size = std::min(new_size * 2, MAX_BUFFER_SIZE);
        need_resize = new_size > m_recv_buffer.size();
    }
    else if (m_recv_buffer_curr_len < new_size * SHRINK_THRESHOLD)
    {
        if (++m_buffer_need_shrink_counter >= SHRINK_DELAY)
        {
            new_size = std::max(new_size / 2, MIN_BUFFER_SIZE);
            need_resize = new_size >= m_recv_buffer_curr_len * 2 && new_size < m_recv_buffer.size();
            if (need_resize)
            {
                m_buffer_need_shrink_counter = 0;
            }
        }
    }
    else
    {
        m_buffer_need_shrink_counter = 0; // 重置收缩计数器
    }

    if (need_resize)
    {
        try
        {
            m_recv_buffer.resize(new_size);
            std::cout << "\t\t\t\t\t\t\t\t\t\t[INFO]: Buffer " << (new_size > m_recv_buffer.size()
                                                                       ? "expanded"
                                                                       : "shrunk")
                << " to " << new_size << " bytes." << std::endl;
        }
        catch (const std::bad_alloc& e)
        {
            std::cerr << "Failed to resize buffer: " << e.what() << std::endl;
            return false;
        }
    }

    return true;
}


// 开始异步写操作
void Session::StartAsyncWrite() {
    std::lock_guard<std::mutex> lock(m_send_mutex);

    // 如果当前没有正在发送的数据且发送队列非空，取出队首Node
    if (!m_current_send_node && !m_send_queue.empty()) {
        m_current_send_node = m_send_queue.front();
        m_send_queue.pop();

        // 序列化Node并获取数据和长度
        auto serialized_data = m_current_send_node->serialize();

        // 开始异步写操作
        auto self(shared_from_this());
        boost::asio::async_write(m_socket,
                                 boost::asio::buffer(serialized_data.buffer.get(), serialized_data.length),
                                 [this, self](const boost::system::error_code& ec, std::size_t bytes_transferred) {
                                     this->HandleWrite(ec, bytes_transferred);
                                 });
    }
}

// 处理写完成的回调
void Session::HandleWrite(const boost::system::error_code& ec, std::size_t bytes_transferred) {
    std::lock_guard<std::mutex> lock(m_send_mutex);

    if (ec) {
        HandleError(ec, "Failed to write data");
        return;
    }

    // 清除当前发送的Node
    m_current_send_node.reset();

    // 检查队列是否还有待发送的数据
    if (!m_send_queue.empty()) {
        StartAsyncWrite();  // 继续发送下一个Node
    }
}

// 示例：向队列添加发送任务
void Session::SendNode(const std::shared_ptr<Node>& node) {
    std::lock_guard<std::mutex> lock(m_send_mutex);
    m_send_queue.push(node);
    if (!m_current_send_node) {
        StartAsyncWrite();  // 如果没有正在发送的数据，立即启动发送
    }
}




void Session::HandleError(const boost::system::error_code& ec, const std::string& message ) const
{
    if (!message.empty())
    {
        std::cerr << "Error: " << message << ": " << ec.message() << std::endl;
    }
    else
    {
        std::cerr << "Error on receive: " << ec.message() << std::endl;
    }
    m_server.ClearSession(this->GetUuid());
}

std::ostream& Session::PrintBuffer(std::ostream& os, const std::vector<std::byte>& buffer, const size_t buffer_size,
                                   const size_t header_size)
{
    constexpr size_t OUTPUT_HEAD_TAIL = 16;

    os << "Buffer size: " << buffer_size << " Payload size: " << buffer_size - header_size << "\n";
    os << "Header (hex): ";
    os << std::hex << std::setfill('0');
    for (size_t i = 0; i < header_size; ++i)
    {
        if (i > 0) os << " ";
        os << std::setw(2) << static_cast<int>(buffer[i]);
    }
    os << "\n";

    os << "Payload (hex): ";
    for (size_t i = header_size; i < buffer_size && i < header_size + OUTPUT_HEAD_TAIL; ++i)
    {
        if (i > header_size) os << " ";
        os << std::setw(2) << static_cast<int>(buffer[i]);
    }
    if (buffer_size > header_size + OUTPUT_HEAD_TAIL) { os << " ..."; } // 如果buffer超过32字节，添加省略号
    for (size_t i = std::max(buffer_size - OUTPUT_HEAD_TAIL, header_size + OUTPUT_HEAD_TAIL); i < buffer_size; ++i)
    {
        if (i > header_size) os << " ";
        os << std::setw(2) << static_cast<int>(buffer[i]);
    }
    os << std::dec << "\n"; // 恢复为十进制输出

    return os;
}

std::optional<Node> Session::BufferProcessor(const std::vector<std::byte>& buffer, const size_t buffer_size)
{
    std::print("Session::bufferProcessor() Processed new message: \n");
    if (buffer_size < Node::MAX_HEADER_SIZE || buffer_size > buffer.size())
    {
        std::cerr << "Invalid message size" << std::endl;
        return std::nullopt;
    }

    try
    {
        Node node(buffer.data(), buffer_size);

        // 在这里可以添加任何额外的处理逻辑
        std::cout << "Print Buffer: \n";
        PrintBuffer(std::cout, buffer, buffer_size);
        std::cout << std::endl;
        std::print("Print RecvNode: \n");
        std::cout << node << std::endl;

        return node;
    }
    catch (const std::exception& e)
    {
        std::cerr << "Error processing buffer: " << e.what() << std::endl;
        return std::nullopt;
    }
}

std::string Session::GetUuid() const
{
    return boost::uuids::to_string(m_uuid);
}

/////////////////////////////

Server::Server(boost::asio::io_context& io_context, const int16_t port)
    : m_io_context(io_context),
      m_port(port),
      m_acceptor(io_context, ip::tcp::endpoint(ip::tcp::v4(), port))
{
    std::cout << "Server start success, listen on port: " << m_port << "\n";

    HandleAccept();
}

Server::~Server()
{
    std::cout << "Clear sessions..." << "\n";
    m_sessions_map.clear();
    std::cout << "Closing server." << "\n";
}

void Server::HandleAccept()
{
    // 这里有坑, 要传入的函数签名的sock是一个右值, 你如果使用左值复制下来, HandleAccept()执行之后, sock将释放, 造成悬空引用.
    m_acceptor.async_accept(
        [this](const boost::system::error_code& ec, ip::tcp::socket&& sock)
        {
            if (!ec)
            {
                try
                {
                    PrintRemoteSockInformation(sock);
                    // 创建session, 将有用的东西放到Session中
                    m_current_session = std::make_shared<Session>(std::move(sock), *this);
                    // 往map中添加session
                    auto [first, second] = m_sessions_map.try_emplace(m_current_session->GetUuid(), m_current_session);

                    m_current_session->StartSession();
                }
                catch (const boost::system::system_error& e)
                {
                    std::cerr << "Error getting remote endpoint: " << e.what() << "\n";
                }
                HandleAccept();
            }
            else
            {
                std::cerr << "Error accepting: " << ec.message() << "\n";
            }
        }
    );
}

void Server::PrintRemoteSockInformation(const ip::tcp::socket& sock)
{
    // 获取对端的 endpoint, 输出对端信息
    const tcp::endpoint remote_endpoint = sock.remote_endpoint();
    const std::string ip_address = remote_endpoint.address().to_string();
    const uint16_t port_number = remote_endpoint.port();
    std::cout << "Accept to: " << ip_address << ":" << port_number << "\n";
}

void Server::ClearSession(const std::string& uuid)
{
    m_sessions_map.erase(uuid);
}

int main(const int argc, char* argv[])
{
    if (argc != 2)
    {
        std::cerr << "Usage: async_tcp_server <port>\n";
        return 1;
    }

    int16_t port;
    // Use std::from_chars to convert the argument to an integer
    const auto result = std::from_chars(argv[1], argv[1] + std::strlen(argv[1]), port);
    if (result.ec != std::errc() || result.ptr != argv[1] + std::strlen(argv[1]))
    {
        std::cerr << "Invalid port number\n";
        return 1;
    }

    try
    {
        boost::asio::io_context io_context;
        Server server(io_context, port);
        io_context.run();
    }
    catch (std::exception& e)
    {
        std::cerr << "Exception: " << e.what() << "\n";
    }

    return 0;
}
